//////////////////////////////////////////////////////////////////
// (c) Copyright 2010-  by Jeongnim Kim
//////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////
//   National Center for Supercomputing Applications &
//   Materials Computation Center
//   University of Illinois, Urbana-Champaign
//   Urbana, IL 61801
//   e-mail: jnkim@ncsa.uiuc.edu
//
// Supported by 
//   National Center for Supercomputing Applications, UIUC
//   Materials Computation Center, UIUC
//////////////////////////////////////////////////////////////////
// -*- C++ -*-
/**@file fft3d_engine_a2a.hpp
 */
#include <fft3d/pencildecomp.hpp>

#ifndef QMCPLUSPLUS_FFT3D_ENGINE_A2A_HPP
#define QMCPLUSPLUS_FFT3D_ENGINE_A2A_HPP
namespace mpi
{

  struct ia2a_dummy
  {
    int posted;
    template<typename T>
      inline void push(communicator& comm, T* data_in, T* data_out, size_t count, int tag)
      {
        MPI_Datatype dtype=get_mpi_datatype<T>(data_in[0]); 
        int n=comm.size();
        if(n==1) return; // nothing to do
        int me=comm.rank();
        posted=0;
        for(int i=0;i<n; ++i)
          if(i!=me) MPI_Irecv(data_out+count*i, count, dtype, i, tag, comm, &(comm.requests[posted++]));
        for(int i=0;i<n; ++i)
          if(i!=me) MPI_Isend(data_in+count*i, count, dtype, i, tag, comm, &(comm.requests[posted++]));
        memcpy(data_out,data_in,sizeof(T)*count);
      }

    inline int collect(communicator& comm)
    {
      return MPI_Waitall(posted,&comm.requests[0],&comm.status[0]);
    }

  };

  struct ia2a_pairexchange
  {
    int posted;

    template<typename T>
      inline void push(communicator& comm, T* data_in, T* data_out, size_t count, int tag)
      {
        MPI_Datatype dtype=get_mpi_datatype<T>(data_in[0]); 
        int n=comm.size();
        if(n==1) return; // nothing to do
        int me=comm.rank();
        posted=0;
        for(int i=1; i<n; ++i)
        {
          int to=(me+i)%n;
          int from=(me-i+n)%n;
          MPI_Irecv(data_out+count*from, count, dtype, from, tag, comm, &(comm.requests[posted++]));
          MPI_Isend(data_in +count*to,   count, dtype,   to, tag, comm, &(comm.requests[posted++]));
        }
        memcpy(data_out,data_in,sizeof(T)*count);
      }

    inline int collect(communicator& comm)
    {
      return MPI_Waitall(posted,&comm.requests[0],&comm.status[0]);
    }
  };

  struct ia2a_a2a
  {
    template<typename T>
      inline void push(communicator& comm, T* data_in, T* data_out, size_t count, int tag)
      {
        alltoall(comm,data_in,data_out,count);
      }

    //always successful
    inline int collect(communicator& comm) 
    {
      return MPI_SUCCESS;
    }
  };

  template<typename T, typename A2AENG>
    struct FFT3DEngineIA2A: public PencilDecomp<T>, public A2AENG
  {
    typedef PencilDecomp<T> LayoutType;

    FFT3DEngineIA2A(int l, int nx, int ny, int ipx, int ipy)
      : LayoutType(l,nx,ny,ipx,ipy)
    { }

    FFT3DEngineIA2A(int l, int ny, communicator& c)
      : LayoutType(l,ny,c)
    {}


    void fft_forward(int nvar)
    {
      qmcplusplus::Timer clock;
      for(int ivar=0; ivar<nvar; ++ivar)
      {
        fft1d(0);
        clock.restart();
        A2AENG::push(*rowComm,&rowBuffer[0],&colBuffer[0],msgSize[IROW],IROW);
        A2AENG::collect(*rowComm);
        wtime[ROWTRANS_T]+=clock.elapsed();

        fft1d(1);
        clock.restart();
        A2AENG::push(*colComm,&colBuffer[0],&rowBuffer[0],msgSize[ICOL],ICOL);
        A2AENG::collect(*colComm);
        wtime[COLTRANS_T]+=clock.elapsed();

        fft1d(2);
      }
    }

    void fft_backward(int nvar)
    {
      qmcplusplus::Timer clock;
      for(int ivar=0; ivar<nvar; ++ivar)
      {
        fft1d(2);
        clock.restart();
        A2AENG::push(*colComm,&colBuffer[0],&rowBuffer[0],msgSize[ICOL],ICOL);
        A2AENG::collect(*colComm);
        wtime[COLTRANS_T]+=clock.elapsed();

        fft1d(1);

        clock.restart();
        A2AENG::push(*rowComm,&rowBuffer[0],&colBuffer[0],msgSize[IROW],IROW);
        A2AENG::collect(*rowComm);
        wtime[ROWTRANS_T]+=clock.elapsed();

        fft1d(0);
      }
    }
  };

//reference
//  template<typename T>
//    struct FFT3DEngineA2A: public PencilDecomp<T>
//  {
//    typedef PencilDecomp<T> LayoutType;
//
//    FFT3DEngineA2A(int l, int nx, int ny, int ipx, int ipy)
//      : LayoutType(l,nx,ny,ipx,ipy)
//    { }
//
//    FFT3DEngineA2A(int l, int ny, communicator& c)
//      : LayoutType(l,ny,c)
//    {}
//
//    void fft_forward(int nvar)
//    {
//      qmcplusplus::Timer clock;
//      for(int ivar=0; ivar<nvar; ++ivar)
//      {
//        fft1d(0);
//        clock.restart();
//        alltoall(*rowComm,&rowBuffer[0],&colBuffer[0],msgSize[IROW]);
//        wtime[ROWTRANS_T]+=clock.elapsed();
//        fft1d(1);
//        clock.restart();
//        alltoall(*colComm,&colBuffer[0],&rowBuffer[0],msgSize[ICOL]);
//        wtime[COLTRANS_T]+=clock.elapsed();
//        fft1d(2);
//      }
//    }
//
//    void fft_backward(int nvar)
//    {
//      qmcplusplus::Timer clock;
//      for(int ivar=0; ivar<nvar; ++ivar)
//      {
//        fft1d(2);
//
//        clock.restart();
//        alltoall(*colComm,&colBuffer[0],&rowBuffer[0],msgSize[ICOL]);
//        wtime[COLTRANS_T]+=clock.elapsed();
//
//        fft1d(1);
//
//        clock.restart();
//        alltoall(*rowComm,&rowBuffer[0],&colBuffer[0],msgSize[IROW]);
//        wtime[ROWTRANS_T]+=clock.elapsed();
//
//        fft1d(0);
//      }
//    }
//  };
//

}
#endif
/***************************************************************************
 * $RCSfile$   $Author: jnkim $
 * $Revision: 1770 $   $Date: 2007-02-17 17:45:38 -0600 (Sat, 17 Feb 2007) $
 * $Id: fft2d.cpp 1770 2007-02-17 23:45:38Z jnkim $ 
 ***************************************************************************/
